#!/bin/bash

# EnPlatform Assistant 本地AI模型管理脚本
# 支持AMD64和ARM64架构，专门优化Excel文件处理功能

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
COMPOSE_FILE="$SCRIPT_DIR/docker-compose.yml"
DEV_COMPOSE_FILE="$SCRIPT_DIR/docker-compose-dev.yml"

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
PURPLE='\033[0;35m'
CYAN='\033[0;36m'
NC='\033[0m' # No Color

# 日志函数
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

log_debug() {
    echo -e "${PURPLE}[DEBUG]${NC} $1"
}

log_step() {
    echo -e "${CYAN}[STEP]${NC} $1"
}

# 检查Docker是否运行
check_docker() {
    if ! docker info >/dev/null 2>&1; then
        log_error "Docker未运行，请先启动Docker"
        exit 1
    fi
    log_debug "Docker运行正常"
}

# 检查系统架构
check_architecture() {
    ARCH=$(uname -m)
    case $ARCH in
        "arm64"|"aarch64")
            log_info "检测到ARM64架构 (Apple Silicon/ARM服务器)"
            export DOCKER_DEFAULT_PLATFORM=linux/arm64
            return 0
            ;;
        "x86_64"|"amd64")
            log_info "检测到AMD64架构 (Intel/AMD处理器)"
            export DOCKER_DEFAULT_PLATFORM=linux/amd64
            return 1
            ;;
        *)
            log_warning "未知架构: $ARCH，默认使用AMD64"
            export DOCKER_DEFAULT_PLATFORM=linux/amd64
            return 1
            ;;
    esac
}

# 检查系统资源
check_system_resources() {
    log_step "检查系统资源..."
    
    # 检查内存
    if command -v free >/dev/null 2>&1; then
        TOTAL_MEM=$(free -g | awk '/^Mem:/{print $2}')
        if [ "$TOTAL_MEM" -lt 8 ]; then
            log_warning "系统内存少于8GB，建议使用轻量级配置"
            LIGHTWEIGHT_MODE=true
        else
            log_success "系统内存充足: ${TOTAL_MEM}GB"
        fi
    elif command -v sysctl >/dev/null 2>&1; then
        # macOS
        TOTAL_MEM_BYTES=$(sysctl -n hw.memsize)
        TOTAL_MEM_GB=$((TOTAL_MEM_BYTES / 1024 / 1024 / 1024))
        if [ "$TOTAL_MEM_GB" -lt 8 ]; then
            log_warning "系统内存少于8GB，建议使用轻量级配置"
            LIGHTWEIGHT_MODE=true
        else
            log_success "系统内存充足: ${TOTAL_MEM_GB}GB"
        fi
    fi
    
    # 检查磁盘空间 (兼容macOS)
    if command -v df >/dev/null 2>&1; then
        if df -h "$SCRIPT_DIR" >/dev/null 2>&1; then
            AVAILABLE_SPACE=$(df -h "$SCRIPT_DIR" | awk 'NR==2 {print $4}')
            log_success "磁盘空间充足: ${AVAILABLE_SPACE}"
        else
            log_info "无法检查磁盘空间"
        fi
    else
        log_info "df命令不可用，跳过磁盘空间检查"
    fi
}

# 创建必要的目录
create_directories() {
    log_step "创建必要的目录..."
    
    mkdir -p "$SCRIPT_DIR/data/ollama"
    mkdir -p "$SCRIPT_DIR/data/english-embeddings"
    mkdir -p "$SCRIPT_DIR/data/chinese-embeddings"
    mkdir -p "$SCRIPT_DIR/config/nginx/conf.d"
    mkdir -p "$SCRIPT_DIR/config/ollama"
    mkdir -p "$SCRIPT_DIR/scripts"
    mkdir -p "$SCRIPT_DIR/logs"
    mkdir -p "$SCRIPT_DIR/test-data"
    
    log_success "目录创建完成"
}

# 选择部署模式
select_deployment_mode() {
    if [ "$LIGHTWEIGHT_MODE" = true ]; then
        log_info "自动选择超轻量级模式（内存不足）"
        DEPLOYMENT_MODE="minimal"
        return
    fi

    echo ""
    echo "请选择部署模式："
    echo "1) 完整模式 (推荐) - 包含所有功能，内存需求8GB+"
    echo "2) 开发模式 - 轻量级配置，快速启动，内存需求4GB+"
    echo "3) 轻量级模式 - 最小配置，仅核心功能，内存需求2GB+"
    echo "4) 超轻量级模式 (NEW) - Excel专用，极低资源，内存需求1.5GB+"
    echo ""
    read -p "请输入选择 (1/2/3/4，默认4): " choice

    case ${choice:-4} in
        1)
            DEPLOYMENT_MODE="full"
            COMPOSE_FILE="$SCRIPT_DIR/docker-compose.yml"
            log_info "选择完整模式"
            ;;
        2)
            DEPLOYMENT_MODE="development"
            COMPOSE_FILE="$DEV_COMPOSE_FILE"
            log_info "选择开发模式"
            ;;
        3)
            DEPLOYMENT_MODE="lightweight"
            COMPOSE_FILE="$SCRIPT_DIR/docker-compose.yml"
            log_info "选择轻量级模式"
            ;;
        4)
            DEPLOYMENT_MODE="minimal"
            COMPOSE_FILE="$SCRIPT_DIR/docker-compose-minimal.yml"
            log_info "选择超轻量级模式（Excel专用）"
            ;;
        *)
            DEPLOYMENT_MODE="minimal"
            COMPOSE_FILE="$SCRIPT_DIR/docker-compose-minimal.yml"
            log_info "默认选择超轻量级模式"
            ;;
    esac
}

# 启动服务
start_services() {
    log_step "启动本地AI服务..."
    
    check_docker
    check_architecture
    check_system_resources
    create_directories
    select_deployment_mode
    
    cd "$SCRIPT_DIR"
    
    # 根据部署模式启动服务
    case $DEPLOYMENT_MODE in
        "full")
            docker-compose -f "$COMPOSE_FILE" up -d
            ;;
        "development")
            docker-compose -f "$DEV_COMPOSE_FILE" up -d
            ;;
        "lightweight")
            docker-compose -f "$COMPOSE_FILE" --profile lightweight up -d
            ;;
        "minimal")
            docker-compose -f "$COMPOSE_FILE" up -d
            ;;
    esac
    
    log_info "等待服务启动..."
    sleep 15
    
    # 检查服务状态
    if check_services_health; then
        log_success "本地AI服务启动成功"
        show_service_info
    else
        log_error "部分服务启动失败，请查看日志"
        show_logs
        exit 1
    fi
}

# 停止服务
stop_services() {
    log_step "停止本地AI服务..."
    cd "$SCRIPT_DIR"
    
    # 尝试停止所有可能的配置
    docker-compose -f "$COMPOSE_FILE" down 2>/dev/null || true
    docker-compose -f "$DEV_COMPOSE_FILE" down 2>/dev/null || true
    
    log_success "本地AI服务已停止"
}

# 重启服务
restart_services() {
    log_step "重启本地AI服务..."
    stop_services
    sleep 3
    start_services
}

# 检查服务健康状态
check_services_health() {
    local all_healthy=true
    
    # 检查Ollama主服务
    if curl -s http://localhost:11434/api/tags >/dev/null 2>&1; then
        log_success "Ollama主服务正常 (端口11434)"
    else
        log_warning "Ollama主服务未响应"
        all_healthy=false
    fi
    
    # 检查英文嵌入服务
    if curl -s http://localhost:8081/api/tags >/dev/null 2>&1; then
        log_success "英文嵌入服务正常 (端口8081)"
    else
        log_warning "英文嵌入服务未响应"
        all_healthy=false
    fi
    
    # 检查中文嵌入服务（仅在完整模式下）
    if [ "$DEPLOYMENT_MODE" = "full" ]; then
        if curl -s http://localhost:8082/api/tags >/dev/null 2>&1; then
            log_success "中文嵌入服务正常 (端口8082)"
        else
            log_warning "中文嵌入服务未响应"
            all_healthy=false
        fi
    fi
    
    return $($all_healthy && echo 0 || echo 1)
}

# 显示服务信息
show_service_info() {
    echo ""
    log_info "=== 服务端点信息 ==="
    echo "  🤖 Ollama主服务 (对话):     http://localhost:11434"
    echo "  🔤 英文嵌入服务:           http://localhost:8081"
    
    if [ "$DEPLOYMENT_MODE" = "full" ]; then
        echo "  🈳 中文嵌入服务:           http://localhost:8082"
        echo "  🌐 API网关:               http://localhost:8080"
    fi
    
    echo ""
    log_info "=== 快速测试命令 ==="
    echo "  ./manage.sh test           # 完整功能测试"
    echo "  ./manage.sh test-excel     # Excel处理测试"
    echo "  ./manage.sh test-chat      # 对话功能测试"
    echo ""
}

# 显示服务状态
show_status() {
    log_step "检查服务状态..."
    cd "$SCRIPT_DIR"
    
    echo ""
    log_info "=== Docker容器状态 ==="
    docker-compose -f "$COMPOSE_FILE" ps 2>/dev/null || docker-compose -f "$DEV_COMPOSE_FILE" ps 2>/dev/null || true
    
    echo ""
    log_info "=== 服务健康检查 ==="
    check_services_health
    
    echo ""
    log_info "=== 系统资源使用 ==="
    docker stats --no-stream --format "table {{.Container}}\t{{.CPUPerc}}\t{{.MemUsage}}" | head -10
}

# 下载推荐模型
download_models() {
    log_step "下载Excel处理优化模型..."

    check_services_health || {
        log_error "服务未运行，请先启动服务"
        exit 1
    }

    # 下载对话模型
    log_info "下载对话模型 Qwen2.5:7b (Excel分析专用)..."
    docker exec enplatform-ollama ollama pull qwen2.5:7b || {
        log_warning "Qwen2.5:7b下载失败，尝试备用模型..."
        docker exec enplatform-ollama ollama pull phi3.5:3.8b
    }

    # 下载英文嵌入模型
    log_info "下载英文嵌入模型 nomic-embed-text..."
    docker exec enplatform-english-embeddings ollama pull nomic-embed-text

    # 下载中文嵌入模型（仅在完整模式下）
    if [ "$DEPLOYMENT_MODE" = "full" ]; then
        log_info "下载中文嵌入模型 bge-m3..."
        docker exec enplatform-chinese-embeddings ollama pull bge-m3 || {
            log_warning "bge-m3下载失败，使用备用模型..."
            docker exec enplatform-chinese-embeddings ollama pull bge-small-zh-v1.5
        }
    fi

    log_success "模型下载完成"
    list_models
}

# 列出已安装模型
list_models() {
    log_step "列出已安装模型..."

    echo ""
    log_info "=== Ollama主服务模型 ==="
    docker exec enplatform-ollama ollama list 2>/dev/null || log_warning "无法获取主服务模型列表"

    echo ""
    log_info "=== 英文嵌入服务模型 ==="
    docker exec enplatform-english-embeddings ollama list 2>/dev/null || log_warning "无法获取英文嵌入模型列表"

    if [ "$DEPLOYMENT_MODE" = "full" ]; then
        echo ""
        log_info "=== 中文嵌入服务模型 ==="
        docker exec enplatform-chinese-embeddings ollama list 2>/dev/null || log_warning "无法获取中文嵌入模型列表"
    fi
}

# 测试对话功能
test_chat() {
    log_step "测试对话功能..."

    echo ""
    log_info "测试Excel数据分析对话..."

    local test_prompt="分析这个Excel数据：属性编码=TEMP_01, 属性名称=环境温度, 数据类型=float, 属性类型=量测属性。请解释这个属性的含义和用途。"

    local response=$(curl -s -X POST http://localhost:11434/api/generate \
        -H "Content-Type: application/json" \
        -d "{
            \"model\": \"qwen2.5:7b\",
            \"prompt\": \"$test_prompt\",
            \"stream\": false,
            \"options\": {
                \"temperature\": 0.7,
                \"top_p\": 0.9
            }
        }" | jq -r '.response' 2>/dev/null)

    if [ -n "$response" ] && [ "$response" != "null" ]; then
        log_success "对话功能测试成功"
        echo "回答: $response"
    else
        log_warning "对话功能测试失败，尝试备用模型..."
        # 尝试备用模型
        local backup_response=$(curl -s -X POST http://localhost:11434/api/generate \
            -H "Content-Type: application/json" \
            -d "{
                \"model\": \"phi3.5:3.8b\",
                \"prompt\": \"$test_prompt\",
                \"stream\": false
            }" | jq -r '.response' 2>/dev/null)

        if [ -n "$backup_response" ] && [ "$backup_response" != "null" ]; then
            log_success "备用模型测试成功"
            echo "回答: $backup_response"
        else
            log_error "对话功能测试失败"
        fi
    fi
}

# 测试嵌入功能
test_embed() {
    log_step "测试嵌入功能..."

    # 测试英文嵌入
    echo ""
    log_info "测试英文嵌入模型..."
    local english_text="Temperature sensor TEMP_01 float measurement attribute"
    local english_result=$(curl -s -X POST http://localhost:8081/api/embeddings \
        -H "Content-Type: application/json" \
        -d "{
            \"model\": \"nomic-embed-text\",
            \"prompt\": \"$english_text\"
        }" | jq '.embedding | length' 2>/dev/null)

    if [ -n "$english_result" ] && [ "$english_result" -gt 0 ]; then
        log_success "英文嵌入测试成功，向量维度: $english_result"
    else
        log_warning "英文嵌入测试失败"
    fi

    # 测试中文嵌入（仅在完整模式下）
    if [ "$DEPLOYMENT_MODE" = "full" ]; then
        echo ""
        log_info "测试中文嵌入模型..."
        local chinese_text="属性编码 环境温度 浮点数 量测属性"
        local chinese_result=$(curl -s -X POST http://localhost:8082/api/embeddings \
            -H "Content-Type: application/json" \
            -d "{
                \"model\": \"bge-m3\",
                \"prompt\": \"$chinese_text\"
            }" | jq '.embedding | length' 2>/dev/null)

        if [ -n "$chinese_result" ] && [ "$chinese_result" -gt 0 ]; then
            log_success "中文嵌入测试成功，向量维度: $chinese_result"
        else
            log_warning "中文嵌入测试失败"
        fi
    fi
}

# Excel处理专项测试
test_excel() {
    log_step "Excel处理功能专项测试..."

    echo ""
    log_info "=== 测试Excel表头解析 ==="
    local headers=("属性编码" "属性名称" "数据类型" "属性类型" "量测类型" "单位")
    for header in "${headers[@]}"; do
        local result=$(curl -s -X POST http://localhost:8081/api/embeddings \
            -H "Content-Type: application/json" \
            -d "{
                \"model\": \"nomic-embed-text\",
                \"prompt\": \"$header\"
            }" | jq '.embedding | length' 2>/dev/null)

        if [ -n "$result" ] && [ "$result" -gt 0 ]; then
            log_success "表头 '$header' 嵌入成功"
        else
            log_warning "表头 '$header' 嵌入失败"
        fi
    done

    echo ""
    log_info "=== 测试Excel数据行处理 ==="
    local data_rows=(
        "TEMP_01 环境温度 float 量测属性 遥测 °C"
        "STATUS_01 设备状态 enum 量测属性 遥信"
        "CTRL_01 控制指令 bool 控制属性"
    )

    for row in "${data_rows[@]}"; do
        local result=$(curl -s -X POST http://localhost:8082/api/embeddings \
            -H "Content-Type: application/json" \
            -d "{
                \"model\": \"bge-m3\",
                \"prompt\": \"$row\"
            }" | jq '.embedding | length' 2>/dev/null)

        if [ -n "$result" ] && [ "$result" -gt 0 ]; then
            log_success "数据行嵌入成功: $(echo "$row" | cut -d' ' -f1-2)"
        else
            log_warning "数据行嵌入失败: $(echo "$row" | cut -d' ' -f1-2)"
        fi
    done

    echo ""
    log_info "=== 测试Excel数据问答 ==="
    local excel_question="这个Excel文件中的TEMP_01属性是什么类型的数据？"
    local answer=$(curl -s -X POST http://localhost:11434/api/generate \
        -H "Content-Type: application/json" \
        -d "{
            \"model\": \"qwen2.5:7b\",
            \"prompt\": \"根据Excel数据：TEMP_01 环境温度 float 量测属性 遥测 °C，回答问题：$excel_question\",
            \"stream\": false
        }" | jq -r '.response' 2>/dev/null)

    if [ -n "$answer" ] && [ "$answer" != "null" ]; then
        log_success "Excel数据问答测试成功"
        echo "问题: $excel_question"
        echo "回答: $answer"
    else
        log_warning "Excel数据问答测试失败"
    fi
}

# 完整功能测试
test_all() {
    log_step "执行完整功能测试..."

    check_services_health || {
        log_error "服务未运行，请先启动服务"
        exit 1
    }

    test_chat
    test_embed
    test_excel

    echo ""
    log_success "=== 测试完成 ==="
    log_info "如果所有测试都通过，说明系统已准备就绪"
    log_info "可以开始使用Excel文件上传和处理功能"
}

# 显示日志
show_logs() {
    local service=${1:-""}
    cd "$SCRIPT_DIR"

    if [ -n "$service" ]; then
        log_info "显示 $service 服务日志..."
        docker-compose -f "$COMPOSE_FILE" logs -f "$service" 2>/dev/null || \
        docker-compose -f "$DEV_COMPOSE_FILE" logs -f "$service" 2>/dev/null || \
        log_error "无法获取服务日志"
    else
        log_info "显示所有服务日志..."
        docker-compose -f "$COMPOSE_FILE" logs -f 2>/dev/null || \
        docker-compose -f "$DEV_COMPOSE_FILE" logs -f 2>/dev/null || \
        log_error "无法获取服务日志"
    fi
}

# 清理数据
clean_data() {
    log_warning "这将删除所有本地AI数据，包括下载的模型"
    read -p "确认继续? (y/N): " -n 1 -r
    echo

    if [[ $REPLY =~ ^[Yy]$ ]]; then
        log_step "清理本地AI数据..."
        cd "$SCRIPT_DIR"

        # 停止服务
        docker-compose -f "$COMPOSE_FILE" down -v 2>/dev/null || true
        docker-compose -f "$DEV_COMPOSE_FILE" down -v 2>/dev/null || true

        # 删除数据目录
        rm -rf "$SCRIPT_DIR/data"
        rm -rf "$SCRIPT_DIR/logs"

        log_success "数据清理完成"
    else
        log_info "取消清理操作"
    fi
}

# 超轻量级模式专用命令
setup_minimal() {
    log_step "设置超轻量级Excel处理模型..."

    # 检查是否有超轻量级模式的容器在运行
    if ! docker ps --format "table {{.Names}}" | grep -q "enplatform-ollama-chat\|enplatform-ollama-embedding"; then
        log_error "请先使用超轻量级模式启动服务"
        log_info "运行: ./manage.sh start 并选择选项4"
        exit 1
    fi

    # 运行超轻量级设置脚本
    ./scripts/minimal-setup.sh setup
}

# 测试超轻量级Excel功能
test_minimal_excel() {
    log_step "测试超轻量级Excel处理功能..."

    # 检查服务是否运行
    if ! docker ps | grep -q "enplatform-ollama-embedding"; then
        log_error "嵌入服务未运行，请先启动服务"
        exit 1
    fi

    # 运行Excel处理测试
    ./scripts/minimal-setup.sh test
}

# 显示资源使用情况
show_resources() {
    log_step "显示系统资源使用情况..."
    ./scripts/minimal-setup.sh resources
}

# 显示帮助信息
show_help() {
    echo "EnPlatform Assistant 本地AI模型管理脚本"
    echo "专门优化Excel文件解析、向量嵌入和对话功能"
    echo ""
    echo "用法: $0 [命令]"
    echo ""
    echo "🚀 基本操作:"
    echo "  start         启动本地AI服务（自动检测架构）"
    echo "  stop          停止本地AI服务"
    echo "  restart       重启服务"
    echo "  status        显示服务状态和健康检查"
    echo ""
    echo "📦 模型管理:"
    echo "  download      下载Excel处理优化模型"
    echo "  setup-minimal 设置超轻量级模型（Excel专用）"
    echo "  models        列出已安装模型"
    echo ""
    echo "🧪 测试功能:"
    echo "  test          完整功能测试"
    echo "  test-chat     测试对话功能"
    echo "  test-embed    测试嵌入功能"
    echo "  test-excel    专门测试Excel处理功能"
    echo "  test-minimal  测试超轻量级Excel功能"
    echo ""
    echo "🛠️ 维护操作:"
    echo "  logs [服务]   查看日志（可选指定服务名）"
    echo "  resources     显示资源使用情况"
    echo "  clean         清理所有数据和模型（谨慎使用）"
    echo "  help          显示此帮助信息"
    echo ""
    echo "📋 快速开始 (超轻量级):"
    echo "  1. $0 start              # 启动服务，选择选项4"
    echo "  2. $0 setup-minimal      # 下载超轻量级模型"
    echo "  3. $0 test-minimal       # 测试Excel功能"
    echo ""
    echo "🏗️ 支持的架构:"
    echo "  ✅ AMD64 (Intel/AMD处理器)"
    echo "  ✅ ARM64 (Apple Silicon M1/M2/M3)"
    echo ""
    echo "💡 部署模式:"
    echo "  1. 完整模式 (8GB+ 内存)"
    echo "  2. 开发模式 (4GB+ 内存)"
    echo "  3. 轻量级模式 (2GB+ 内存)"
    echo "  4. 超轻量级模式 (1.5GB+ 内存) - Excel专用"
    echo ""
    echo "📚 相关文档:"
    echo "  README.md                # 详细部署指南"
    echo "  docs/                    # 完整文档目录"
}

# 主函数
main() {
    case "${1:-help}" in
        start)
            start_services
            ;;
        stop)
            stop_services
            ;;
        restart)
            restart_services
            ;;
        status)
            show_status
            ;;
        download)
            download_models
            ;;
        setup-minimal)
            setup_minimal
            ;;
        models)
            list_models
            ;;
        test)
            test_all
            ;;
        test-chat)
            test_chat
            ;;
        test-embed)
            test_embed
            ;;
        test-excel)
            test_excel
            ;;
        test-minimal)
            test_minimal_excel
            ;;
        logs)
            show_logs "$2"
            ;;
        resources)
            show_resources
            ;;
        clean)
            clean_data
            ;;
        help|--help|-h)
            show_help
            ;;
        *)
            log_error "未知命令: $1"
            show_help
            exit 1
            ;;
    esac
}

main "$@"
